# -*- coding:utf-8 -*-
# author: lkz
# date: 2016/08/05 15:25

import os
import time
import socket
import struct
import signal
import random

from multiprocessing import Queue
from Queue import Empty as QueueEmpty
from icmp import ICMPSocket
from packet_receiver import PacketReceiver
from ip_icmp_parser import ICMP_DATA_PARSER, ICMP_HEADER_PARSER, SOCKET_BUF_SIZE

import argparse

DEFAULT_MIN_RTT = 10000


class Ping(object):
    def __init__(self, targets, period=15, packets_once=10, run_forever=False, data_process_func=None):
        # ping参数
        self.targets = targets
        self.period = period
        self.packets_once = packets_once
        self.__packets_once = float(packets_once)

        # ping 数据处理函数
        self.__data_process_func = data_process_func

        # 程序控制参数
        self.__run_forever = run_forever
        self.__sequence = 0
        self.__pid = os.getpid()
        self.__invalid_addr_list = []
        self.__ip_list = []

        # 子程序, 与子程序程序之间交互的队列
        self.__packet_queue = Queue()
        self.__packet_receiver = None

        # ICMP Socket
        self.__icmp = ICMPSocket(8, 0, ICMP_DATA_PARSER.size)

        # ping数据统计相关变量
        self.__ip_statistics = {
            # ip: [packet_loss, max_rtt, min_rtt, average_rtt]
        }

        # 不能解析的地址的统计数据(只初始化一次), 现在没有处理这种情形
        self.__invalid_addr_statistics = {}

    def start(self):
        for target in self.targets:
            try:
                ip_addr = socket.gethostbyname(target)
                self.__ip_list.append(ip_addr)
            except socket.gaierror:
                self.__invalid_addr_list.append(target)
                self.__invalid_addr_statistics[target] = [1, -1, -1, -1]

        self.__packet_receiver = PacketReceiver(
            self.__packet_queue,
            self.__ip_list,
            SOCKET_BUF_SIZE
        )

        self.__packet_receiver.daemon = True
        self.__packet_receiver.start()

        if not self.__run_forever:
            self.__ping()
            self.__receive_and_process_data()
            return

        while self.__run_forever:
            start = time.time()

            if not self.__packet_receiver.is_alive():
                self.stop()
                break

            self.__ping()
            self.__receive_and_process_data()

            time_consumed = time.time() - start
            sleep_time = self.period - time_consumed

            if sleep_time < 0:
                print "oh my god"
                # self.stop()
                # return

            time.sleep(sleep_time)

    def stop(self):
        self.__icmp and self.__icmp.close()
        self.__icmp = None

        self.__packet_receiver.stop()
        self.__packet_queue.close()

        self.__run_forever = False
        os.kill(self.__pid, signal.SIGKILL)

    def __get_inc_sequence(self):
        ret = self.__sequence
        self.__sequence += 1

        if self.__sequence == 2147483647:
            self.__sequence = 0

        return ret

    def __ping(self):
        sequence = self.__get_inc_sequence()

        for i in range(self.packets_once):
            for ip in self.__ip_list:
                time.sleep(random.uniform(0.0009, 0.001))
                msg = time.time()
                data = ICMP_DATA_PARSER.pack(self.__pid, sequence, msg)
                self.__icmp and self.__icmp.sendto(ip, data)

    def __receive_and_process_data(self):
        """
        接收并处理packet_parser进程发送到data_queue的数据, 这个函数尽量不再调用别的函数以提高运行速度
        :return:
        """
        # 从包队列中接收ip包
        start_time = time.time()

        raw_data = []
        while True:
            try:
                dst_ip, end_time, raw_ip_data = self.__packet_queue.get_nowait()
                raw_data.append((dst_ip, end_time, raw_ip_data))
            except QueueEmpty as e:
                time_used = time.time() - start_time

                # 发送完所有的icmp包后, 按道理, 所有活的ip都会在4秒内返回数据
                if time_used > 4:
                    break

        # 解析数据
        data = []
        for dst_ip, end_time, raw_ip_data in raw_data:
            icmp_header_fields = ICMP_HEADER_PARSER.unpack(raw_ip_data[20:24])

            # Just receive packet of icmp type is 0 and icmp code is 0
            if not (icmp_header_fields[0] == 0 and icmp_header_fields[1] == 0):
                continue

            try:
                icmp_data_fields = ICMP_DATA_PARSER.unpack(raw_ip_data[24:])
            except:
                continue

            pid = icmp_data_fields[0]
            if pid != self.__pid:
                continue

            rtt = (end_time - icmp_data_fields[2]) * 1000

            data.append((dst_ip, rtt))

        # 初始化ip数据的统计结构
        for ip in self.__ip_list:
            self.__ip_statistics[ip] = [0, 0, DEFAULT_MIN_RTT, 0]

        # 统计
        for ip, rtt in data:
            statistics = self.__ip_statistics[ip]

            # 更新packet_num
            statistics[0] += 1

            # 更新max rtt
            if rtt > statistics[1]:
                statistics[1] = rtt

            # 更新min rtt
            if rtt < statistics[2]:
                statistics[2] = rtt

            # 更新sum rtt
            statistics[3] += rtt

        # 最终统计结果
        for _, statistics in self.__ip_statistics.items():
            statistics[0] = 1 - statistics[0] / self.__packets_once

            if statistics[1] == 0:
                statistics[1] = -1

            if statistics[2] == DEFAULT_MIN_RTT:
                statistics[2] = -1

            if statistics[3] == 0:
                statistics[3] = -1
            else:
                statistics[3] /= self.__packets_once

        if self.__data_process_func:
            self.__data_process_func(self.__ip_statistics)
            return

        self.__default_data_process_func(self.__ip_statistics)

    def __default_data_process_func(self, data):
        out_tpl = (
            "{ip}, packet loss {packet_loss:.2f}, "
            "max rtt {max_rtt:.3f}, min rtt {min_rtt:.3f}, average rtt {avg_rtt:.3f}"
        )

        for ip, statistics in data.items():

            if statistics[1] <= 0:
                continue

            print out_tpl.format(
                ip=ip,
                packet_loss=statistics[0],
                max_rtt=statistics[1],
                min_rtt=statistics[2],
                avg_rtt=statistics[3]
            )

        for ip, statistics in data.items():
            if statistics[1] > 0:
                continue

            print ip, "timeout"

        for ip in self.__invalid_addr_list:
            print ip, "unknown host %s" % ip


def parse_args():
    parser = argparse.ArgumentParser(description="fast ping", version="1.0")

    ip_group = parser.add_mutually_exclusive_group(required=True)
    ip_group.add_argument("-i", "--ip", dest="ip",  nargs="+", help="ip or dns")
    ip_group.add_argument("-f", "--ip-file", dest="ip_file", help="file containing ip list line by line")
    ip_group.add_argument("-n", "--net", dest="net", help="net_addr like 192.168.1.1/24")

    return parser.parse_args()


def _parse_net_addr(net_addr_and_mask):
    str_net_addr, mask = net_addr_and_mask.split("/")
    ip_parser = struct.Struct("!I")

    mask = int(mask)
    try:
        int_net_addr = ip_parser.unpack(socket.inet_aton(str_net_addr))[0]
    except:
        print "Invalid net address: %s" % str_net_addr
        import sys
        sys.exit(0)

    ip_count = (2 ** (32 - mask)) - 2
    return [socket.inet_ntoa(ip_parser.pack(int_net_addr + i)) for i in range(1, ip_count + 1)]


if __name__ == "__main__":
    args = parse_args()
    ip_list = []

    if args.ip:
        ip_list.extend(args.ip)
    elif args.ip_file:
        with open(args.ip_file, mode="r") as f:
            ip_list = map(lambda x: x.strip(),  f.readlines())
            ip_list = filter(lambda x: bool(x), ip_list)
    else:
        ip_list = _parse_net_addr(args.net)

    ping = Ping(ip_list, packets_once=5, run_forever=False)
    ping.start()